import numpy as np
import ufl
from mpi4py import MPI
from petsc4py.PETSc import ScalarType

from dolfinx import mesh, fem
from dolfinx.fem import functionspace, Function, Constant, dirichletbc, form, assemble_scalar
from dolfinx.fem.petsc import LinearProblem
from ufl import grad, div, dot, dx, ds, TrialFunction, TestFunction, inner

# 创建网格，区域为 [-1, 1]×[-1, 1]
# nx, ny = 32, 32
nx, ny = 64, 64
domain = mesh.create_rectangle(MPI.COMM_WORLD, [np.array([-1.0, -1.0]), np.array([1.0, 1.0])],
                              [nx, ny], mesh.CellType.triangle)

# 定义函数空间，使用 Lagrange 一次元
V = functionspace(domain, ("Lagrange", 1))

# 定义解析解
def u_exact(x):
    return np.exp(x[0] + x[1])

# 定义变分问题中的 TrialFunction 和 TestFunction
u = TrialFunction(V)
v = TestFunction(V)

# 定义方程的右端项
f = fem.Function(V)
f.interpolate(lambda x: -2 * np.exp(x[0] + x[1]))

# 定义边界条件
# Neumann边界不需要放到bcs 直接在变分形式中处理
#  Dirichlet 边界条件：x=-1, x=1, y=1
# 标记 x=-1 的边界
boundary_facets_x_neg = mesh.locate_entities_boundary(
    domain, domain.topology.dim - 1,
    lambda x: np.isclose(x[0], -1.0)
)
dofs_x_neg = fem.locate_dofs_topological(V, domain.topology.dim - 1, boundary_facets_x_neg)
u_bc_x_neg = fem.Function(V)

u_bc_x_neg.interpolate(lambda x: np.exp(-1 + x[1]))

bc_x_neg = dirichletbc(u_bc_x_neg, dofs_x_neg)

# 标记 x=1 的边界
boundary_facets_x_pos = mesh.locate_entities_boundary(
    domain, domain.topology.dim - 1,
    lambda x: np.isclose(x[0], 1.0)
)
dofs_x_pos = fem.locate_dofs_topological(V, domain.topology.dim - 1, boundary_facets_x_pos)
u_bc_x_pos = fem.Function(V)
u_bc_x_pos.interpolate(lambda x: np.exp(1 + x[1]))
bc_x_pos = dirichletbc(u_bc_x_pos, dofs_x_pos)

# 标记 y=1 的边界
boundary_facets_y_pos = mesh.locate_entities_boundary(
    domain, domain.topology.dim - 1,
    lambda x: np.isclose(x[1], 1.0)
)
dofs_y_pos = fem.locate_dofs_topological(V, domain.topology.dim - 1, boundary_facets_y_pos)
u_bc_y_pos = fem.Function(V)
u_bc_y_pos.interpolate(lambda x: np.exp(x[0] + 1))
bc_y_pos = dirichletbc(u_bc_y_pos, dofs_y_pos)


#  Neumann 边界条件：y=-1
boundary_facets_y_neg = mesh.locate_entities_boundary(
    domain, domain.topology.dim - 1,
    lambda x: np.isclose(x[1], -1.0)
)
# 构建 Neumann 边界条件对应的线性形式贡献
g_neumann = fem.Function(V)
g_neumann.interpolate(lambda x: -np.exp(x[0] - 1))
L_neumann = g_neumann * v * ds



# 构建变分问题的双线性形式和线性形式
a = dot(grad(u), grad(v)) * dx
L = inner(f, v) * dx + L_neumann

# 汇总 Dirichlet 边界条件
bcs = [bc_x_neg, bc_x_pos, bc_y_pos]

# 创建线性问题并求解
problem = LinearProblem(a, L, bcs=bcs, petsc_options={"ksp_type": "preonly", "pc_type": "lu"})
u_sol = problem.solve()

# 计算误差
u_ex_func = fem.Function(V)
u_ex_func.interpolate(u_exact)
e = u_sol - u_ex_func

# 计算 L2 误差
L2_error = np.sqrt(domain.comm.allreduce(assemble_scalar(form(dot(e, e) * dx)), op=MPI.SUM))

# 计算 H1 半范数误差
H1_semi_error = np.sqrt(domain.comm.allreduce(assemble_scalar(form(dot(grad(e), grad(e)) * dx)), op=MPI.SUM))

# 主进程输出误差信息
if domain.comm.rank == 0:
    print(f"L2 误差: {L2_error:.6e}")
    print(f"H1 半范数误差: {H1_semi_error:.6e}")